import math

import torch

from algorithms.algorithm import Algorithm
from individual.population import Population
from operators.crossover.sbx import SBX
from operators.mutations.polynomial_mutation import PolynomialMutation
from problems.problem import Problem
from utils.vectors import get_ref_vectors


class MOEAD(Algorithm):
    parameters = [
        {
            'label': 'type', 'name': 'type', 'type': 'select', 'description': '聚合函数', 'defaultValue': 4,
            "options": [
                {'label': "PBI approach", "value": 1},
                {'label': "Tchebycheff approach", "value": 2},
                {'label': "Tchebycheff norm", "value": 3},
                {'label': "Modified Tchebycheff", "value": 4}]
        }
    ]

    def __init__(self, problem: Problem,
                 pop_size: int,
                 kwarg):
        self.ref_vectors = get_ref_vectors(pop_size, problem.obj_dim)
        pop_size = len(self.ref_vectors)
        super().__init__(pop_size, problem, kwarg)
        self.population = self.problem.init_population(pop_size)
        self.ideal_point = self.population.obj.min(0)[0].unsqueeze(dim=0)
        self.neighbors_size = math.ceil(pop_size / 10)
        self.neighbors = torch.cdist(self.ref_vectors, self.ref_vectors).sort(1)[1][:, :self.neighbors_size]

    def each_iteration(self):
        dec = self.population.dec
        obj = self.population.obj
        for i in range(self.pop_size):
            parent = self.neighbors[i, torch.randperm(self.neighbors_size)]
            offspring_dec = SBX.do(dec[parent[0]], dec[parent[1]], self.problem, offspring_size=1)
            offspring_dec = PolynomialMutation.do(offspring_dec, self.problem)
            offspring = self.problem.estimate_population(offspring_dec)
            offspring_obj = offspring.obj
            self.ideal_point = torch.min(torch.cat((self.ideal_point, offspring_obj), dim=0), dim=0)[0].unsqueeze(dim=0)
            if self.type == 1:
                # PBI approach
                normW = torch.sqrt(torch.sum(self.ref_vectors[parent] ** 2, dim=1))
                normP = torch.sqrt(
                    torch.sum((obj[parent] - self.ideal_point.repeat(self.neighbors_size, 1)) ** 2, dim=1))
                normO = torch.sqrt(torch.sum((offspring_obj - self.ideal_point) ** 2, dim=1))
                CosineP = torch.sum(
                    (obj[parent] - self.ideal_point.repeat(self.neighbors_size, 1)) * self.ref_vectors[parent],
                    dim=1) / normW / normP
                CosineO = torch.sum(
                    (offspring_obj - self.ideal_point).repeat(self.neighbors_size, 1) * self.ref_vectors[parent],
                    dim=1) / normW / normO
                g_old = normP * CosineP + 5 * normP * torch.sqrt(1 - CosineP ** 2)
                g_new = normO * CosineO + 5 * normO * torch.sqrt(1 - CosineO ** 2)

            elif self.type == 2:
                # Tchebycheff approach
                g_old = torch.max(
                    torch.abs(obj[parent] - self.ideal_point.repeat(self.neighbors_size, 1)) * self.ref_vectors[parent],
                    dim=1)[0]
                g_new = torch.max(
                    (torch.abs(offspring_obj - self.ideal_point)).repeat(self.neighbors_size, 1) * self.ref_vectors[
                        parent], dim=1)[0]

            elif self.type == 3:
                # Tchebycheff approach with normalization
                max_point = torch.max(self.population.obj, dim=0)[0]
                g_old = torch.max(torch.abs(obj[parent] - self.ideal_point.repeat(self.neighbors_size, 1)) / (
                        max_point - self.ideal_point).repeat(self.neighbors_size, 1) * self.ref_vectors[parent], dim=1)[
                    0]
                g_new = torch.max(torch.abs(offspring_obj - self.ideal_point) / (max_point - self.ideal_point).repeat(
                    self.neighbors_size, 1) * self.ref_vectors[parent], dim=1)[0]

            elif self.type == 4:
                # Modified Tchebycheff approach
                g_old = torch.max(
                    torch.abs(obj[parent] - self.ideal_point.repeat(self.neighbors_size, 1)) / self.ref_vectors[parent],
                    dim=1)[0]
                g_new = torch.max(
                    torch.abs(offspring_obj - self.ideal_point).repeat(self.neighbors_size, 1) / self.ref_vectors[
                        parent], dim=1)[0]
            else:
                raise ValueError(f"不支持类型为 {self.type} 的聚合函数")
            dec[parent[g_old >= g_new]] = offspring_dec
            obj[parent[g_old >= g_new]] = offspring_obj
        self.population = Population(dec, obj)
